#include "ulib.h"
#include "elf.h"

#include "syscall.h"

#define debug(x ...) // sc_kprintf(x)
#define error(x ...) sc_kprintf(x)
#define PAGE_SIZE 4096

/* Some constants */
static const char * EH_TYPE[5] =
{
	 "NONE",
	 "Relocatable",
	 "Executable",
	 "Shared Object",
	 "Core",
};

static const char * EH_MACHINE[8] =
{
	 "NONE",
	 "AT&T WE 32100",
	 "SPARC",
	 "Intel 80386",
	 "Motorola 68000",
	 "Motorola 88000",
	 "Intel 80860",
	 "MIPS RS3000",
};

static const char * SH_TYPE[12] =
{
	 "NULL",
	 "PROGBITS",
	 "SYMTAB",
	 "STRTAB",
	 "RELA",
	 "HASH",
	 "DYNAMIC",
	 "NOTE",
	 "NOBITS",
	 "REL",
	 "SHLIB",
	 "DYNSYM",
};

#define SH_TYPE_SYMTAB 2
#define SH_TYPE_STRTAB 3
#define SH_TYPE_NOBITS 8
#define SH_TYPE_RELOC  9

static const char * SYMTAB_BIND[3] =
{
	 "LOCAL",
	 "GLOBAL",
	 "WEAK",
};

static const char * SYMTAB_TYPE[5] =
{
	 "NO TYPE",
	 "OBJECT",
	 "FUNCTION",
	 "SECTION",
	 "FILE",  
};

#define GET_RELOC_SYM(i)  ((i)>>8)
#define GET_RELOC_TYPE(i) ((unsigned char)(i))

#define GET_SYMTAB_BIND(i)   ((i)>>4)
#define GET_SYMTAB_TYPE(i)   ((i)&0xf)

#define SHN_UNDEF  0
#define SHN_COMMON 0xfff2

#define MOD_LOAD_SYMBOL   "module_load"
#define MOD_UNLOAD_SYMBOL "module_unload"

static const char *
get_symbol_string(void *elf, uint32_t index)
{  
	 uint32_t i;
	 const char * ret;
	 struct elf_t *eh;
	 struct secthdr_t *sh;
	 struct symtab_t *symtab;

	 eh = (struct elf_t*)elf;

	 for (i = 0; i < eh->e_shnum; i ++)
	 {  
		  sh = (struct secthdr_t*)((uintptr_t)elf + eh->e_shoff + (i * eh->e_shentsize));
		  if(sh->sh_type == SH_TYPE_SYMTAB)
		  {
			   symtab = (struct symtab_t*)((uintptr_t)elf + sh->sh_offset);
			   sh = (struct secthdr_t*)(elf + eh->e_shoff + (sh->sh_link * eh->e_shentsize));
        
			   if(sh->sh_type == SH_TYPE_STRTAB)
			   {
					if (index == 0)
					{
						 return " ";
					}
					else return (char *)((uintptr_t)elf + sh->sh_offset + index);
			   }
		  }
	 }

	 ret = "";
	 return ret;
}

#define CHARS_MAX            10240
static uint32_t char_count;
static char     chars[CHARS_MAX];

#define EXPORT_SYM_COUNT_MAX 1024
#define EXPORT_SYM_HASH      197

static uint32_t    ex_sym_count;

static int         ex_sym_f[EXPORT_SYM_HASH];
static const char *ex_sym_name[EXPORT_SYM_COUNT_MAX];
static uintptr_t   ex_sym_ptr[EXPORT_SYM_COUNT_MAX];
static uint32_t    ex_sym_flags[EXPORT_SYM_COUNT_MAX];
static int         ex_sym_n[EXPORT_SYM_COUNT_MAX];

static uint32_t
sym_hash(const char *name, uint32_t len)
{
	 int idx;
	 uint32_t result = 0;
	 for (idx = 0; idx != len; ++ idx)
		  result = (result * 13 + name[idx]) % EXPORT_SYM_HASH;
	 return result;
}


static void
touch_export_sym(const char *name, uintptr_t ptr, uint32_t flags)
{
	 int name_len = strlen(name);
	 int h = sym_hash(name, name_len);
	 int cur = ex_sym_f[h];

	 while (cur != -1)
	 {
		  if (strcmp(ex_sym_name[cur], name) == 0)
		  {
			   break;
		  }
		  else
		  {
			   cur = ex_sym_n[cur];
		  }
	 }

	 if (cur == -1)
	 {
		  cur = ex_sym_count ++;
		  ex_sym_n[cur] = ex_sym_f[h];
		  ex_sym_f[h] = cur;

		  char *_name = chars + char_count;
		  memmove(_name, name, name_len);
		  char_count += name_len;
		  ex_sym_name[cur] = _name;
		  chars[char_count ++] = 0;
	 }
	 
	 ex_sym_ptr[cur] = ptr;
	 ex_sym_flags[cur] = flags;
}

static int
find_export_sym(const char *name, int touch)
{
	 int name_len = strlen(name);
	 int h = sym_hash(name, name_len);
	 int cur = ex_sym_f[h];

	 while (cur != -1)
	 {
		  if (strcmp(ex_sym_name[cur], name) == 0)
		  {
			   break;
		  }
		  else
		  {
			   cur = ex_sym_n[cur];
		  }
	 }

	 if (cur == -1 && touch)
	 {
		  cur = ex_sym_count ++;
		  ex_sym_n[cur] = ex_sym_f[h];
		  ex_sym_f[h] = cur;

		  char *_name = chars + char_count;
		  memmove(_name, name, name_len);
		  char_count += name_len;
		  chars[char_count ++] = 0;

		  ex_sym_name[cur] = _name;
	 }

	 return cur;
}

#define MOD_COUNT_MAX 1024

static uint32_t    mod_count;

struct mod_status_t
{
	 int       valid;
	 int       rkv_id;
	 uintptr_t ptr;
	 uint32_t  size;
	 uintptr_t common_ptr;
	 uint32_t  common_size;
	 uintptr_t load_ptr;
	 uintptr_t unload_ptr;
};

static int         mod_f[EXPORT_SYM_HASH];
static const char *mod_name[MOD_COUNT_MAX];
static int         mod_n[MOD_COUNT_MAX];

static struct mod_status_t
mod_status[MOD_COUNT_MAX];


static int
find_mod(const char *name, int touch)
{
	 int name_len = strlen(name);
	 int h = sym_hash(name, name_len);
	 int cur = mod_f[h];

	 while (cur != -1)
	 {
		  if (strcmp(mod_name[cur], name) == 0)
		  {
			   break;
		  }
		  else
		  {
			   cur = mod_n[cur];
		  }
	 }

	 if (cur == -1 && touch)
	 {
		  cur = mod_count ++;
		  mod_n[cur] = mod_f[h];
		  mod_f[h] = cur;

		  char *_name = chars + char_count;
		  memmove(_name, name, name_len);
		  char_count += name_len;
		  chars[char_count ++] = 0;

		  mod_name[cur] = _name;
	 }

	 return cur;
}

struct symtab_t *
fill_symbol_struct(void *elf, uint32_t symbol)
{
	 uint32_t i;
	 struct elf_t     *eh;
	 struct secthdr_t *sh;
	 struct symtab_t  *symtab;

	 eh = (struct elf_t*)elf;

	 for(i = 0; i < eh->e_shnum; i++)
	 {  
		  sh = (struct secthdr_t*)(elf + eh->e_shoff + (i * eh->e_shentsize));
		  if (sh->sh_type == SH_TYPE_SYMTAB)
		  {
			   symtab = (struct symtab_t*)((uintptr_t)elf + sh->sh_offset + (symbol * sh->sh_entsize));
			   return (struct symtab_t *)symtab;
		  }
	 }

	 return (struct symtab_t *)0;
}

uintptr_t
get_section_offset(void *elf, uint32_t info)
{
	 struct elf_t     *eh;
	 struct secthdr_t *sh;

	 eh = (struct elf_t*)elf;
	 sh = (struct secthdr_t*)(elf + eh->e_shoff + (info * eh->e_shentsize));

	 uintptr_t result;
	 if (sh->sh_type == SH_TYPE_NOBITS)
		  result = sh->sh_addr;
	 else result = sh->sh_offset;
	 return result;
}

static int
elf_header_check(void *elf)
{
	 struct elf_t *eh = (struct elf_t *)elf;
	 if (eh->e_magic != ELF_MAGIC)
	 {
		  // invalid signature
		  error("invalid signature\n");
		  return -1;
	 }

	 if(eh->e_type != 0x01)
	 {
		  error("error type = %d\n", eh->e_type);
		  // error type
		  return -1;
	 }

	 if(eh->e_machine != 0x03)
	 {
		  error("error machine type %d\n", eh->e_machine);
		  // error machine type
		  return -1;
	 }

	 if(eh->e_entry != 0x0)
	 {
		  error("error entry point %08x\n", eh->e_entry);
		  return -1;
	 }

	 return 0;
}

static int
elf_mod_parse(void *elf, const char *name, int export_symbol,
			  uintptr_t *common_data, uint32_t *common_size,
			  uintptr_t *mod_load_ptr, uintptr_t *mod_unload_ptr)
{
	 uint32_t i, x;
	 uintptr_t reloc_addr, mem_addr;
	 uintptr_t mod_load = 0, mod_unload = 0;
		
	 struct elf_t     *eh;
	 struct secthdr_t *sh;
	 struct reloc_t   *reloc;
	 struct symtab_t  *symtab;

	 uintptr_t cur_common_alloc = 0;
	 uintptr_t cur_common_align = 1;

	 eh = (struct elf_t*)elf;

	 /* do section header stuff here */
	 for (i = 0; i < eh->e_shnum; i++)
	 {  
		  sh = (struct secthdr_t*)(elf + eh->e_shoff + (i * eh->e_shentsize));		  
		  if (sh->sh_type == SH_TYPE_SYMTAB)
		  {
			   for(x = 0; x < sh->sh_size; x += sh->sh_entsize)
			   {
					symtab = (struct symtab_t*)(elf + sh->sh_offset + x);

					debug("found sym: [%s] info[%02x] size[%d] addr[%08x]",
						  get_symbol_string(elf, symtab->sym_name),
						  symtab->sym_info,
						  symtab->sym_size,
						  symtab->sym_address);
					if (symtab->sym_shndx != SHN_UNDEF && symtab->sym_shndx < 0xff00)
					{
						 debug(" value offset [%08x]", 
							   get_section_offset(elf, symtab->sym_shndx) +
							   symtab->sym_address);

						 switch (GET_SYMTAB_BIND(symtab->sym_info))
						 {
						 case 0:
							  debug(" local\n");
							  // local symbol
							  if (strcmp(get_symbol_string(elf, symtab->sym_name),
										 MOD_LOAD_SYMBOL) == 0)
							  {
								   mod_load = get_section_offset(elf, symtab->sym_shndx) +
										symtab->sym_address + (uint32_t)elf;
							  }
							  else if (strcmp(get_symbol_string(elf, symtab->sym_name),
											  MOD_UNLOAD_SYMBOL) == 0)
							  {
								   mod_unload = get_section_offset(elf, symtab->sym_shndx) +
										symtab->sym_address + (uint32_t)elf;
							  }
							  break;
						 case 1:
							  debug(" global\n");
							  // global
							  if (export_symbol)
								   elf_mod_touch_symbol(get_symbol_string(elf, symtab->sym_name),
														(void *)(symtab->sym_address + (uint32_t)elf), 0);
							  break;
						 case 2:
							  debug(" weak\n");
							  // weak
							  if (export_symbol)
								   elf_mod_create_symbol(get_symbol_string(elf, symtab->sym_name),
														 (void *)(symtab->sym_address + (uint32_t)elf), 0);
							  break;
						 }

					}
					else if (symtab->sym_shndx == SHN_COMMON)
					{
						 debug(" alloc %d byte align 0x%x\n",
							   symtab->sym_size, symtab->sym_address);
						 if (bsf(symtab->sym_address) != bsr(symtab->sym_address))
						 {
							  error(" bad align\n");
							  return -1;
						 }
						 if (symtab->sym_address > cur_common_align)
							  cur_common_align = symtab->sym_address;
						 cur_common_alloc = ((cur_common_alloc - 1) | (symtab->sym_address - 1)) + 1;
						 symtab->sym_address = cur_common_alloc;
						 cur_common_alloc += symtab->sym_size;
					}
					else
					{
						 debug(" shndx[%04x]\n", symtab->sym_shndx);
					}
			   }
		  }
		  else if (sh->sh_type == SH_TYPE_NOBITS)
		  {
			   debug(" alloc %d byte align 0x%x\n",
					 sh->sh_size, sh->sh_addralign);
			   if (bsf(sh->sh_addralign) != bsr(sh->sh_addralign))
			   {
					error(" bad align\n");
					return -1;
			   }
			   if (sh->sh_addralign > cur_common_align)
					cur_common_align = sh->sh_addralign;
			   cur_common_alloc = ((cur_common_alloc - 1) | (sh->sh_addralign - 1)) + 1;
			   sh->sh_addr = cur_common_alloc;
			   cur_common_alloc += sh->sh_size;
		  }
	 }

	 debug("common: alloc %d bytes, align 0x%x\n",
		   cur_common_alloc,
		   cur_common_align);

	 uintptr_t common_space;
	 if (cur_common_align > PAGE_SIZE)
	 {
		  error("align failed\n");
		  return -1;
	 }
	 else if (cur_common_alloc > 0)
	 {
		  common_space = (uintptr_t)bmalloc(cur_common_alloc);
		  debug("common %08x, size %d\n", common_space, cur_common_alloc);
		  
		  *common_data = common_space;
		  *common_size = cur_common_alloc;
	 }
	 else
	 {
		  debug("no common\n");
		  *common_data = 0;
		  *common_size = 0;
	 }

	 if (mod_load == 0 || mod_unload == 0)
	 {
		  error("entry point was not found\n");
		  return -1;
	 }
	 else
	 {
		  *mod_load_ptr = mod_load;
		  *mod_unload_ptr = mod_unload;
	 }
        		
	 /* fill the relocation entries */
	 for (i = 0; i < eh->e_shnum; i++)
	 {  
		  sh = (struct secthdr_t*)(elf + eh->e_shoff + (i * eh->e_shentsize));

		  if (sh->sh_type == SH_TYPE_RELOC)
		  {
			   for (x = 0; x < sh->sh_size; x += sh->sh_entsize)
			   {
					reloc = (struct reloc_t*)(elf + sh->sh_offset + x);
					symtab = fill_symbol_struct(elf, GET_RELOC_SYM(reloc->rl_info));

					debug("reloc[%02x] offset[%08x] for [%s], sym offset[%08x]\n",
						  GET_RELOC_TYPE(reloc->rl_info),
						  reloc->rl_offset,
						  get_symbol_string(elf, symtab->sym_name),
						  symtab->sym_address
						 );

					mem_addr = (uintptr_t)elf + reloc->rl_offset;
					mem_addr += get_section_offset(elf, sh->sh_info);
                
					/* external reference (kernel symbol most likely) */
					if(symtab->sym_shndx == SHN_UNDEF)
					{
						 int idx = find_export_sym(get_symbol_string(elf, symtab->sym_name), 0);
						 if (idx == -1)
						 {
							  if (strcmp(get_symbol_string(elf, symtab->sym_name), name) != 0)
							  {
								   error("unresolved dependency \"%s\", set with 0\n",
										 get_symbol_string(elf, symtab->sym_name));
								   // goto failed_with_common_space;
								   reloc_addr = 0;
							  }
							  else reloc_addr = (uintptr_t)elf;
						 }
						 else
						 {
							  debug("extern symbol addr = %08x\n", ex_sym_ptr[idx]);
							  reloc_addr = ex_sym_ptr[idx];
						 }
					}
					else if (symtab->sym_shndx < 0xff00)
					{
						 debug("section offset %08x, addr %08x\n", get_section_offset(elf, symtab->sym_shndx), symtab->sym_address);
						 if (((struct secthdr_t*)((uint32_t)eh + eh->e_shoff + (symtab->sym_shndx * eh->e_shentsize)))->sh_type == SH_TYPE_NOBITS)
							  reloc_addr = common_space;
						 else reloc_addr = (uint32_t)elf;
						 reloc_addr += get_section_offset(elf, symtab->sym_shndx);
						 reloc_addr += symtab->sym_address;
					}
					else if (symtab->sym_shndx == SHN_COMMON)
					{
						 reloc_addr = common_space + symtab->sym_address;
					}

					debug("old value = %08x\n", *(uintptr_t *)mem_addr);

					switch (GET_RELOC_TYPE(reloc->rl_info))
					{
					case 0x01:
						 *(uintptr_t *)mem_addr = reloc_addr + *(uintptr_t *)mem_addr;
						 debug("fill abs address %08x to %08x\n", *(uintptr_t *)mem_addr, mem_addr);
						 break;
					case 0x02:
						 reloc_addr = reloc_addr - mem_addr;
						 *(uintptr_t *)mem_addr = reloc_addr + *(uintptr_t *)mem_addr ;
						 debug("fill rel address %08x to %08x\n", *(uintptr_t *)mem_addr, mem_addr);
						 break;
					default:
						 error("unsupported relocation type (%x)\n", GET_RELOC_TYPE(reloc->rl_info));
						 break;
					}
			   }
		  }
	 }
	 return 0;
	 
failed_with_common_space:
	 if (cur_common_alloc > 0)
		  bfree((void *)common_space);
	 return -1;
}

int
elf_mod_init(void)
{
	 ex_sym_count = 0;
	 char_count = 0;
	 mod_count = 0;
	 
	 memset(ex_sym_f, -1, sizeof(ex_sym_f));
	 memset(mod_f, -1, sizeof(mod_f));
	 
	 return 0;
}

int
elf_mod_set_symbol(const char *name, void *ptr, uint32_t flags)
{
	 int idx = find_export_sym(name, 0);
	 if (idx == -1) return -1;

	 ex_sym_ptr[idx] = (uintptr_t)ptr;
	 ex_sym_flags[idx] = flags;
	 
	 return 0;
}

void
elf_mod_touch_symbol(const char *name, void *ptr, uint32_t flags)
{
	 int idx = find_export_sym(name, 1);

	 ex_sym_ptr[idx] = (uintptr_t)ptr;
	 ex_sym_flags[idx] = flags;	 
}

int
elf_mod_create_symbol(const char *name, void *ptr, uint32_t flags)
{
	 int idx = find_export_sym(name, 1);
	 if (idx != ex_sym_count - 1) return -1;

	 ex_sym_ptr[idx] = (uintptr_t)ptr;
	 ex_sym_flags[idx] = flags;
	 
	 return 0;
}

int
elf_mod_get_symbol(const char *name, void **ptr, uint32_t *flags)
{
	 int idx = find_export_sym(name, 0);
	 if (idx == -1) return -1;
	 
	 *ptr = (void *)ex_sym_ptr[idx];
	 *flags = ex_sym_flags[idx];
}

int
elf_mod_load(const char *name)
{
	 debug("loading mod %s\n", name);
	 
	 int idx = find_mod(name, 0);
	 if ((idx != -1) && (mod_status[idx].valid))
		  return 0;

	 if (idx == -1) return -1;

	 mod_status[idx].valid = 0;
	 int mod_id = 0;

	 const char *ptr;
	 if (rkv_get(mod_status[idx].rkv_id, &ptr, &mod_status[idx].size) != 0)
		  return -1;

	 if (elf_header_check((void *)ptr))
		  return -1;

	 struct elf_t *eh = (struct elf_t *)ptr;
	 if ((mod_status[idx].ptr = (uintptr_t)bmalloc(mod_status[idx].size)) == 0)
		  return -1;

	 debug("mod address %08x - %08x %08x\n", mod_status[idx].ptr, mod_status[idx].ptr + mod_status[idx].size, mod_status[idx].size);
	 memmove((void *)mod_status[idx].ptr, ptr, mod_status[idx].size);

	 debug("parsing\n");
	 if (elf_mod_parse((void *)mod_status[idx].ptr, name, 0,
					   &mod_status[idx].common_ptr, &mod_status[idx].common_size,
					   &mod_status[idx].load_ptr, &mod_status[idx].unload_ptr) != 0)
	 {
		  bfree((void *)mod_status[idx].ptr);
		  return -1;
	 }

	 sc_kprintf("mod %s: %08x size %08x common %08x\n",
				name, mod_status[idx].ptr,
				mod_status[idx].size, mod_status[idx].common_size);

	 debug("load() begin\n");
	 if (((int(*)(void))(mod_status[idx].load_ptr))() != 0)
	 {
		  if (mod_status[idx].common_size != 0)
			   bfree((void *)mod_status[idx].common_ptr);
		  bfree((void *)mod_status[idx].ptr);
		  error("load() failed\n");
		  return -1;
	 }
	 debug("load() success\n");

	 mod_status[idx].valid = 1;
	 return 0;
}

void
elf_mod_register(const char *name, int rkv_id)
{
	 debug("register for mod %s\n", name);
	 int idx = find_mod(name, 1);
	 if ((idx != mod_count - 1))
		  return;

	 mod_status[idx].valid = 0;
	 mod_status[idx].rkv_id = rkv_id;
}
